﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;

namespace g4
{
    /// <summary>
    /// 网格局部参数化类：计算网格上局部区域的UV映射（参数化）
    /// 支持多种参数化模式，基于图遍历（类Dijkstra算法）从种子区域向外传播UV坐标
    /// 适用于生成网格局部区域的连续参数化坐标
    /// </summary>
    public class MeshLocalParam
    {
        /// <summary>
        /// 无效UV坐标标记（用于表示未计算或超出范围的UV）
        /// </summary>
        public static readonly Vector2f InvalidUV = new Vector2f(float.MaxValue, float.MaxValue);


        /// <summary>
        /// UV参数化模式枚举
        /// </summary>
        public enum UVModes
        {
            ExponentialMap,          // 指数映射：基于局部坐标系传播UV，保留角度信息
            ExponentialMap_UpwindAvg, // 迎风平均指数映射：融合多个邻居的UV信息，提高平滑性
            PlanarProjection          // 平面投影：直接将3D点投影到种子帧的XY平面
        }

        /// <summary>
        /// 选择使用的UV参数化模式
        /// </summary>
        public UVModes UVMode = UVModes.ExponentialMap_UpwindAvg;


        /// <summary>
        /// 图节点类：用于图遍历的节点，存储顶点ID、父节点、距离、UV坐标等信息
        /// 继承DynamicPriorityQueueNode以支持优先队列，实现IEquatable用于比较
        /// </summary>
        class GraphNode : DynamicPriorityQueueNode, IEquatable<GraphNode>
        {
            public int id;               // 对应网格顶点的ID
            public GraphNode parent;     // 父节点（图遍历中的前驱）
            public float graph_distance; // 从种子点到该节点的图距离（沿网格表面的累积距离）
            public Vector2f uv;          // 计算得到的UV坐标
            public bool frozen;          // 标记是否已完成UV计算（冻结状态）

            /// <summary>
            /// 节点相等性判断（基于顶点ID）
            /// </summary>
            public bool Equals(GraphNode other)
            {
                return id == other.id;
            }
        }

        /// <summary>
        /// 稀疏优先队列：用于图遍历的优先级队列，按graph_distance排序
        /// </summary>
        DynamicPriorityQueue<GraphNode> SparseQueue;

        /// <summary>
        /// 稀疏节点列表：存储所有参与计算的图节点（以顶点ID为索引）
        /// </summary>
        SparseObjectList<GraphNode> SparseNodes;

        /// <summary>
        /// 节点内存池：用于高效创建和回收GraphNode对象
        /// </summary>
        MemoryPool<GraphNode> SparseNodePool;

        /// <summary>
        /// 获取顶点位置的函数（输入顶点ID，输出3D位置）
        /// </summary>
        Func<int, Vector3f> PositionF;

        /// <summary>
        /// 获取顶点法线的函数（输入顶点ID，输出3D法线）
        /// </summary>
        Func<int, Vector3f> NormalF;

        /// <summary>
        /// 获取顶点邻居的函数（输入顶点ID，输出邻居顶点ID集合）
        /// </summary>
        Func<int, IEnumerable<int>> NeighboursF;

        /// <summary>
        /// 种子帧：参数化的初始坐标系（包含原点和X/Y/Z轴，定义局部参考系）
        /// </summary>
        Frame3f SeedFrame;

        /// <summary>
        /// 最大图距离：计算过程中遇到的最大graph_distance（沿网格表面的距离）
        /// </summary>
        float max_graph_distance;

        /// <summary>
        /// 最大UV距离：计算得到的UV坐标中，距原点的最大距离
        /// </summary>
        float max_uv_distance;


        /// <summary>
        /// 构造函数：初始化网格局部参数化器
        /// </summary>
        /// <param name="nMaxID">最大顶点ID（用于初始化稀疏数据结构）</param>
        /// <param name="nodePositionF">获取顶点位置的函数</param>
        /// <param name="nodeNormalF">获取顶点法线的函数</param>
        /// <param name="neighboursF">获取顶点邻居的函数</param>
        public MeshLocalParam(int nMaxID,
            Func<int, Vector3f> nodePositionF,
            Func<int, Vector3f> nodeNormalF,
            Func<int, IEnumerable<int>> neighboursF)
        {
            PositionF = nodePositionF;
            NormalF = nodeNormalF;
            NeighboursF = neighboursF;

            SparseQueue = new DynamicPriorityQueue<GraphNode>();
            SparseNodes = new SparseObjectList<GraphNode>(nMaxID, 0);
            SparseNodePool = new MemoryPool<GraphNode>();
            max_graph_distance = float.MinValue;
            max_uv_distance = float.MinValue;
        }


        /// <summary>
        /// 重置计算状态：清空队列、节点列表和内存池
        /// </summary>
        public void Reset()
        {
            SparseQueue.Clear(false);
            SparseNodes.Clear();
            SparseNodePool.ReturnAll();
            max_graph_distance = float.MinValue;
        }


        /// <summary>
        /// 核心计算方法：从种子区域开始，计算距离不超过fMaxGraphDistance的顶点的UV坐标
        /// 采用类Dijkstra算法，按图距离优先遍历并传播UV
        /// </summary>
        /// <param name="seedFrame">种子坐标系（定义参数化的初始参考系）</param>
        /// <param name="seedNbrs">种子邻居顶点ID（参数化的起始顶点）</param>
        /// <param name="fMaxGraphDistance">最大图距离（超过此距离的顶点不计算）</param>
        public void ComputeToMaxDistance(Frame3f seedFrame, Index3i seedNbrs, float fMaxGraphDistance)
        {
            SeedFrame = seedFrame;

            // 初始化种子邻居顶点的UV（作为计算起点）
            for (int j = 0; j < 3; ++j)
            {
                int vid = seedNbrs[j];
                GraphNode g = get_node(vid);
                // 计算种子顶点在种子帧中的初始UV
                g.uv = compute_local_uv(ref SeedFrame, PositionF(vid));
                // 图距离初始化为UV的长度（近似沿表面距离）
                g.graph_distance = g.uv.Length;
                g.frozen = true; // 种子顶点的UV固定
                Debug.Assert(!SparseQueue.Contains(g));
                SparseQueue.Enqueue(g, g.graph_distance); // 加入优先队列
            }

            // 图遍历主循环：按图距离从小到大处理节点
            while (SparseQueue.Count > 0)
            {
                GraphNode current = SparseQueue.Dequeue(); // 取出距离最小的节点
                max_graph_distance = Math.Max(current.graph_distance, max_graph_distance);

                // 若当前节点距离超过最大值，终止计算
                if (max_graph_distance > fMaxGraphDistance)
                    return;

                // 根据选择的模式更新当前节点的UV坐标
                if (current.parent != null)
                {
                    switch (UVMode)
                    {
                        case UVModes.ExponentialMap:
                            update_uv_expmap(current);
                            break;
                        case UVModes.ExponentialMap_UpwindAvg:
                            update_uv_upwind_expmap(current);
                            break;
                        case UVModes.PlanarProjection:
                            update_uv_planar(current);
                            break;
                    }
                }

                // 更新最大UV距离
                float uvDistSqr = current.uv.LengthSquared;
                if (uvDistSqr > max_uv_distance)
                    max_uv_distance = uvDistSqr;

                current.frozen = true; // 标记为已处理（UV固定）
                update_neighbours_sparse(current); // 处理当前节点的邻居
            }

            // 计算最终的最大UV距离（开平方）
            max_uv_distance = (float)Math.Sqrt(max_uv_distance);
        }


        /// <summary>
        /// 变换所有计算得到的UV坐标（缩放和平移）
        /// </summary>
        /// <param name="fScale">缩放因子</param>
        /// <param name="vTranslate">平移向量</param>
        public void TransformUV(float fScale, Vector2f vTranslate)
        {
            foreach (var pair in SparseNodes.NonZeroValues())
            {
                GraphNode g = pair.Value;
                if (g.frozen) // 只处理已冻结（计算完成）的节点
                    g.uv = (g.uv * fScale) + vTranslate;
            }
        }


        /// <summary>
        /// 获取计算过程中遇到的最大图距离
        /// </summary>
        public float MaxGraphDistance
        {
            get { return max_graph_distance; }
        }

        /// <summary>
        /// 获取计算得到的最大UV距离
        /// </summary>
        public float MaxUVDistance
        {
            get { return max_uv_distance; }
        }

        /// <summary>
        /// 获取指定顶点的UV坐标
        /// </summary>
        /// <param name="id">顶点ID</param>
        /// <returns>UV坐标（未计算则返回InvalidUV）</returns>
        public Vector2f GetUV(int id)
        {
            GraphNode g = SparseNodes[id];
            return (g == null) ? InvalidUV : g.uv;
        }


        /// <summary>
        /// 将计算得到的UV坐标应用到网格（通过回调函数）
        /// </summary>
        /// <param name="applyF">应用UV的回调函数（输入顶点ID和UV坐标）</param>
        public void ApplyUVs(Action<int, Vector2f> applyF)
        {
            foreach (var pair in SparseNodes.NonZeroValues())
            {
                GraphNode g = pair.Value;
                if (g.frozen) // 只应用已冻结（计算完成）的UV
                    applyF(g.id, g.uv);
            }
        }


        /// <summary>
        /// 计算3D点在指定坐标系（帧）中的局部UV坐标（投影到XY平面）
        /// </summary>
        /// <param name="f">参考坐标系（帧）</param>
        /// <param name="pos">3D点坐标</param>
        /// <returns>2D UV坐标（X为点在f.X轴的投影，Y为在f.Y轴的投影）</returns>
        Vector2f compute_local_uv(ref Frame3f f, Vector3f pos)
        {
            pos -= f.Origin; // 将点转换到帧的局部坐标系
            return new Vector2f(pos.Dot(f.X), pos.Dot(f.Y)); // 投影到XY轴
        }


        /// <summary>
        /// 传播UV坐标：基于邻居的UV和局部坐标系变换，计算当前点的UV
        /// 用于指数映射模式，考虑邻居坐标系与种子坐标系的旋转关系
        /// </summary>
        /// <param name="pos">当前点3D坐标</param>
        /// <param name="nbrUV">邻居点的UV坐标</param>
        /// <param name="fNbr">邻居点的局部坐标系</param>
        /// <param name="fSeed">种子坐标系</param>
        /// <returns>当前点的UV坐标</returns>
        Vector2f propagate_uv(Vector3f pos, Vector2f nbrUV, ref Frame3f fNbr, ref Frame3f fSeed)
        {
            // 计算当前点在邻居坐标系中的局部UV
            Vector2f localUv = compute_local_uv(ref fNbr, pos);

            // 调整种子坐标系，使其Z轴与邻居坐标系Z轴对齐（保持法线一致）
            Frame3f seedToLocal = fSeed;
            seedToLocal.AlignAxis(2, fNbr.Z);

            // 计算邻居坐标系X轴与调整后种子坐标系X轴的夹角
            Vector3f alignedSeedX = seedToLocal.X;
            Vector3f localX = fNbr.X;
            float cosTheta = localX.Dot(alignedSeedX);

            // 计算旋转角度的正弦值（考虑符号）
            float tmp = 1 - cosTheta * cosTheta;
            tmp = Math.Max(tmp, 0); // 避免负数值（浮点误差）
            float sinTheta = (float)Math.Sqrt(tmp);
            Vector3f cross = localX.Cross(alignedSeedX);
            if (cross.Dot(fNbr.Z) < 0) // 根据叉积方向调整符号
                sinTheta = -sinTheta;

            // 构建旋转矩阵（将邻居局部UV旋转到种子坐标系下）
            Matrix2f frameRotate = new Matrix2f(
                cosTheta, sinTheta,
                -sinTheta, cosTheta
            );

            // 邻居UV加上旋转后的局部UV，得到当前点UV
            return nbrUV + frameRotate * localUv;
        }


        /// <summary>
        /// 用指数映射更新节点的UV坐标
        /// 基于父节点的UV和坐标系变换计算当前节点UV
        /// </summary>
        void update_uv_expmap(GraphNode node)
        {
            int vid = node.id;
            Util.gDevAssert(node.parent != null && node.parent.frozen); // 确保父节点已处理
            int parentId = node.parent.id;

            // 父节点的位置和局部坐标系
            Vector3f parentPos = PositionF(parentId);
            Frame3f parentFrame = new Frame3f(parentPos, NormalF(parentId));
            // 传播父节点的UV到当前节点
            node.uv = propagate_uv(PositionF(vid), node.parent.uv, ref parentFrame, ref SeedFrame);
        }


        /// <summary>
        /// 用迎风平均指数映射更新节点的UV坐标
        /// 融合所有已处理邻居的UV信息（加权平均），提高平滑性
        /// </summary>
        void update_uv_upwind_expmap(GraphNode node)
        {
            int vid = node.id;
            Vector3f pos = PositionF(vid);

            Vector2f avgUv = Vector2f.Zero;
            float weightSum = 0;
            int nbrCount = 0;

            // 遍历所有邻居，收集已处理邻居的UV
            foreach (var nbrId in NeighboursF(node.id))
            {
                GraphNode nbrNode = get_node(nbrId, false);
                if (nbrNode.frozen)
                { // 只考虑已冻结（处理完成）的邻居
                    Vector3f nbrPos = PositionF(nbrId);
                    Frame3f nbrFrame = new Frame3f(nbrPos, NormalF(nbrId));
                    // 计算从该邻居传播到当前节点的UV
                    Vector2f nbrUv = propagate_uv(pos, nbrNode.uv, ref nbrFrame, ref SeedFrame);
                    // 权重与距离平方成反比（近邻权重更高）
                    float weight = 1.0f / (pos.DistanceSquared(nbrPos) + MathUtil.ZeroTolerancef);
                    avgUv += weight * nbrUv;
                    weightSum += weight;
                    nbrCount++;
                }
            }

            Util.gDevAssert(nbrCount > 0); // 确保至少有一个已处理的邻居

            // 加权平均得到当前节点的UV
            avgUv /= weightSum;
            node.uv = avgUv;
        }


        /// <summary>
        /// 用平面投影更新节点的UV坐标
        /// 直接将3D点投影到种子坐标系的XY平面
        /// </summary>
        void update_uv_planar(GraphNode g)
        {
            g.uv = compute_local_uv(ref SeedFrame, PositionF(g.id));
        }


        /// <summary>
        /// 获取或创建图节点（从内存池或稀疏列表）
        /// </summary>
        /// <param name="id">顶点ID</param>
        /// <param name="bCreateIfMissing">若不存在是否创建</param>
        /// <returns>图节点（若不创建且不存在则返回null）</returns>
        GraphNode get_node(int id, bool bCreateIfMissing = true)
        {
            GraphNode g = SparseNodes[id];
            if (g == null && bCreateIfMissing)
            {
                g = SparseNodePool.Allocate(); // 从内存池获取节点
                g.id = id;
                g.parent = null;
                g.frozen = false;
                g.uv = Vector2f.Zero;
                g.graph_distance = float.MaxValue; // 初始距离设为最大值
                SparseNodes[id] = g;
            }
            return g;
        }


        /// <summary>
        /// 更新当前节点的邻居节点：计算距离并加入优先队列
        /// 用于图遍历的扩展步骤
        /// </summary>
        void update_neighbours_sparse(GraphNode parent)
        {
            Vector3f parentPos = PositionF(parent.id);
            float parentDist = parent.graph_distance;

            // 遍历所有邻居
            foreach (int nbrId in NeighboursF(parent.id))
            {
                GraphNode nbr = get_node(nbrId);
                if (nbr.frozen) // 跳过已处理的邻居
                    continue;

                // 计算邻居节点的图距离（父节点距离 + 父节点到邻居的欧氏距离）
                float nbrDist = parentDist + parentPos.Distance(PositionF(nbrId));

                // 更新邻居节点的距离和父节点（若更优）
                if (SparseQueue.Contains(nbr))
                {
                    if (nbrDist < nbr.priority)
                    { // 找到更短的路径
                        nbr.parent = parent;
                        nbr.graph_distance = nbrDist;
                        SparseQueue.Update(nbr, nbrDist); // 更新队列优先级
                    }
                }
                else
                {
                    // 新节点加入队列
                    nbr.parent = parent;
                    nbr.graph_distance = nbrDist;
                    SparseQueue.Enqueue(nbr, nbrDist);
                }
            }
        }
    }
}